import os
import re
import sys
import cv2 as cv
from python_ai.common.xcommon import *
from PyCmpltrtok.common_opencv import *
from PyCmpltrtok.common import *
import argparse

if '__main__' == __name__:

    def _main():
        parser = argparse.ArgumentParser()
        parser.add_argument('-n', '--num', help='How many pictures to check at the same time.', type=int, default=2)
        parser.add_argument('-w', '--width', help='How wide your screen view is.', type=int, default=1500)
        parser.add_argument('--height', help='How high your screen view is.', type=int, default=700)
        parser.add_argument('-o', '--offset', help='From which offset to start.', type=int, default=0)
        parser.add_argument('-d', '--dir', help='The directory of labels.', type=str, default=None)

        NUM = 2
        W, H = 1500, 700
        xpos = 0

        args = parser.parse_args()
        NUM = args.num
        W, H = args.width, args.height
        xpos = args.offset
        DIR = args.dir

        TGT_RECT = (W // NUM, H)

        # category names
        label_names_path = '/var/asuspei/my_svn/c_darknet/darknet-debug-cpu/data/voc.names'
        with open(label_names_path, 'r') as f:
            names = f.readlines()
        names = [name[:-2] if name[-2] == '\r' else name[:-1] for name in names]
        # print('\n'.join(names))
        # sys.exit(0)

        # dir of labels and images
        label_dir = '/var/asuspei/my_svn/c_darknet/darknet-debug-cpu/scripts/VOCdevkit/VOC2007/labels'
        if DIR is not None:
            label_dir = DIR
        img_dir = os.path.join(label_dir, '..', 'JPEGImages')

        # map main name to path and to boxes
        id_list = []
        id2path, id2label_path = {}, {}
        for base_name in os.listdir(label_dir):
            xid, _ = os.path.splitext(base_name)
            id_list.append(xid)
            img_path = os.path.join(img_dir, f'{xid}.jpg')
            id2path[xid] = img_path
            id2label_path[xid] = os.path.join(label_dir, base_name)
        id_list = sorted(id_list)

        # read in box info
        def get_boxes(xid):
            label_path = id2label_path[xid]
            with open(label_path, 'r') as f:
                lines = f.readlines()
            boxes = []
            lines = [line[:-2] if line[-2] == '\r' else line[:-1] for line in lines]
            for line in lines:
                vals = line.split()
                name = names[int(vals[0])]
                x, y, w, h = [float(val) for val in vals[1:]]
                box_info = [name]
                box_info.extend([x - w/2, y - h/2, x + w/2, y + h/2])
                boxes.append(box_info)
            return boxes

        xlen = len(id_list)
        while True:
            xids = id_list[xpos:xpos+NUM]
            img_arr = []
            for xid in xids:
                xpath = id2path[xid]
                xboxes = get_boxes(xid)
                img = cv.imread(xpath, cv.IMREAD_COLOR)
                h, w = img.shape[:2]
                for box in xboxes:
                    label = box[0]
                    xy12 = np.array(box[1:]) * np.array([w, h, w, h])
                    xy12 = np.round(xy12).astype(int)
                    x1, y1, x2, y2 = xy12
                    color = rand_color()
                    cv.rectangle(img, (x1, y1), (x2, y2), color, 2)
                    cv.putText(img, f'{label}', (x1, y1 + 15), cv.FONT_HERSHEY_PLAIN, 1, color, 2)
                img = imzoom_rect(img, TGT_RECT)
                color = rand_color()
                cv.putText(img, xid, (0, 15), cv.FONT_HERSHEY_PLAIN, 1, color, 2)
                img_arr.append(img)
            img = np.concatenate(img_arr, axis=1)
            cv.line(img, (TGT_RECT[0], 1), TGT_RECT, (0, 255, 0), 2)
            cv.imshow('check YOLO valid', img)
            k = cv.waitKey()
            if ord('p') == k:
                xpos -= 1
            elif ord('n') == k:
                xpos += 1
            elif ord('q') == k:
                print('Bye!')
                break
            xpos = xpos if xpos >= 0 else 0
            xpos = xpos if xpos < xlen else xlen - 1
            cv.destroyAllWindows()

    _main()
